import argparse
from ultralytics import YOLO
from pathlib import Path
import yaml
import os
from CommunicationClass import CommunicationBase

#Qt主程序调用方法：
#1：切换到环境python.exe目录，运行命令python E:\ultralytics11xu\train.py即可
#2：将算法包拷贝到Lib\site-packages目录下，切换到环境python.exe目录，运行命令python Lib\site-packages\ultralytics11xu\train.py即可
#调用命令，带参数
#python Lib\site-packages\ultralytics11xu\train.py --project_dir="E:/data/" --class_names="class1,class2,class3"
#BaseTrainer中完整参数
"""
    A base class for creating trainers.

    Attributes:
        args (SimpleNamespace): Configuration for the trainer.
        validator (BaseValidator): Validator instance.
        model (nn.Module): Model instance.
        callbacks (defaultdict): Dictionary of callbacks.
        save_dir (Path): Directory to save results.
        wdir (Path): Directory to save weights.
        last (Path): Path to the last checkpoint.
        best (Path): Path to the best checkpoint.
        save_period (int): Save checkpoint every x epochs (disabled if < 1).
        batch_size (int): Batch size for training.
        epochs (int): Number of epochs to train for.
        start_epoch (int): Starting epoch for training.
        device (torch.device): Device to use for training.
        amp (bool): Flag to enable AMP (Automatic Mixed Precision).
        scaler (amp.GradScaler): Gradient scaler for AMP.
        data (str): Path to data.
        trainset (torch.utils.data.Dataset): Training dataset.
        testset (torch.utils.data.Dataset): Testing dataset.
        ema (nn.Module): EMA (Exponential Moving Average) of the model.
        resume (bool): Resume training from a checkpoint.
        lf (nn.Module): Loss function.
        scheduler (torch.optim.lr_scheduler._LRScheduler): Learning rate scheduler.
        best_fitness (float): The best fitness value achieved.
        fitness (float): Current fitness value.
        loss (float): Current loss value.
        tloss (float): Total loss value.
        loss_names (list): List of loss names.
        csv (Path): Path to results CSV file.
    """

def parse_args():
    parser = argparse.ArgumentParser(description='Train a segmentor')
    #             data (str): Path to dataset configuration file.
    #             epochs (int): Number of training epochs.
    #             batch_size (int): Batch size for training.
    #             imgsz (int): Input image size.
    #             device (str): Device to run training on (e.g., 'cuda', 'cpu').
    #             workers (int): Number of worker threads for data loading.
    #             optimizer (str): Optimizer to use for training.
    #             lr0 (float): Initial learning rate.
    #             patience (int): Epochs to wait for no observable improvement for early stopping of training.
    parser.add_argument('--uuid',type=str,default="uuid",help='Project uuid,used to communication')
    parser.add_argument('--model',type=str,default="yolo11s.pt",help='Model configuration file (.yaml or .pt)')
    # parser.add_argument('--model',type=str,default="yolo11s_PromptRestormer.yaml",help='Model configuration file (.yaml or .pt)')
    parser.add_argument('--data',type=str,default="data.yaml",help='Path to dataset configuration file')
    parser.add_argument('--epoch',type=int,default=100,help='Number of training epochs')
    parser.add_argument('--batch_size',type=int,default=4,help='Batch size for training')
    parser.add_argument('--imgsz',type=int,default=640,help='Input image size')
    parser.add_argument('--device',type=str,default="cuda",help="Device to run training on (e.g., 'cuda', 'cpu')")
    parser.add_argument('--workers',type=int,default=1,help='Number of worker threads for data loading')
    parser.add_argument('--optimizer',type=str,help="Optimizer to use for training")
    parser.add_argument('--lr0',type=float,help='Initial learning rate')
    parser.add_argument('--patience',type=int,default=10000,help='Epochs to wait for no observable improvement for early stopping of training')
    parser.add_argument('--project_dir',type=str,default="D:/reduan",help="Project dir")
    parser.add_argument('--class_names',type=str,default='class1',help="Class names,(e.g., 'class1,class2,class3')")
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    # print(args)
    comBase=CommunicationBase(args.uuid)
    # while True:
    #     rst=com.getMessage()
    #     if rst!='':
    #         if rst=="Stop":
    #             print("接收到退出命令！")
    #             break 
    # print(args.project_dir,str(args.class_names))
    try:
        #根据项目生成训练数据配置data.yaml
        project_dir=args.project_dir
        currentPath=Path(__file__).parent.resolve()
        dataYamlPath=str(currentPath.joinpath('ultralytics','cfg','datasets','data.yaml'))
        classList=args.class_names.split(',')
        data={}
        data['path']=os.path.join(project_dir,'temp')
        data['train']='images/train'
        data['val']='images/val'
        data['test']=None
        names={index:onename for index,onename in enumerate(classList)}
        data['names']=names
        #若val文件夹下为空，则验证集和训练集一样
        if len(os.listdir(os.path.join(data['path'],data['val'])))==0:
            data['val']='images/train'
        with open(dataYamlPath, 'w') as file:
            yaml.safe_dump(data, file)

        #model保存在此目录
        modelPath=str(currentPath.joinpath(args.model))
        if modelPath.endswith('.pt'):
            args.model=str(currentPath.joinpath(args.model))
        # Load a model
        model = YOLO(args.model)
        # model = YOLO("yolo11n.yaml")
        # Train the model，传递通讯类实例
        train_results = model.train(
            data=args.data,  # path to dataset YAML
            epochs=args.epoch,  # number of training epochs
            batch=args.batch_size,
            imgsz=args.imgsz,  # training image size
            workers=args.workers,
            device=0,  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
            # save_dir='D:/customRun',
            # wdir='D:/customRun',
            # last='customRun',
            # best='customRun',
            # csv='customRun'
            patience=args.epoch,
            name=os.path.join(project_dir,'saves'),   #训练信息保存目录
            com=comBase
        )

        # Evaluate model performance on the validation set
        # metrics = model.val()

        # Perform object detection on an image
        # results = model("path/to/image.jpg")
        # results[0].show()

        comBase.sendMessage("Info:ExportModel")
        # Export the model to ONNX format
        path = model.export(format="onnx")  # return path to exported model
    except Exception as e:  # 捕获所有类型的异常
        print("Error:"+str(e))
        comBase.sendMessage("Error:"+str(e))
    finally:
        comBase.sendMessage("Finish")
        comBase.destroy()